#include "auto_navigation/mapping/ndt_mapper.hpp"

namespace auto_navigation {

NDTMapper::NDTMapper(ros::NodeHandle& nh, ros::NodeHandle& private_nh)
    : nh_(nh)
    , private_nh_(private_nh)
    , tf_listener_(tf_buffer_)
    , map_cloud_(new pcl::PointCloud<pcl::PointXYZ>)
    , current_cloud_(new pcl::PointCloud<pcl::PointXYZ>)
    , is_first_frame_(true)
{
    initializeParams();
    
    // 初始化订阅器
    points_sub_ = nh_.subscribe("velodyne_points", 1, &NDTMapper::pointCloudCallback, this);
    
    // 初始化发布器
    map_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("ndt/map", 1);
    current_pose_pub_ = nh_.advertise<geometry_msgs::PoseStamped>("ndt/current_pose", 1);
    current_scan_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("ndt/current_scan", 1);
    
    // 初始化服务
    save_map_service_ = nh_.advertiseService("save_map", &NDTMapper::saveMapCallback, this);
    
    // 初始化NDT参数
    ndt_.setTransformationEpsilon(ndt_epsilon_);
    ndt_.setStepSize(ndt_step_size_);
    ndt_.setResolution(ndt_resolution_);
    ndt_.setMaximumIterations(max_iterations_);
    
    // 初始化位姿
    current_pose_ = Eigen::Matrix4f::Identity();
    previous_pose_ = Eigen::Matrix4f::Identity();
    
    ROS_INFO("NDT Mapper initialized");
}

NDTMapper::~NDTMapper() {
    std::string map_file = ros::package::getPath("auto_navigation") + "/maps/map.pcd";
    if (saveMap(map_file)) {
        ROS_INFO("Map saved automatically on shutdown to %s", map_file.c_str());
    }
}

void NDTMapper::initializeParams() {
    // 从参数服务器读取参数，如果没有则使用默认值
    private_nh_.param("voxel_leaf_size", voxel_leaf_size_, 0.3);
    private_nh_.param("ndt_resolution", ndt_resolution_, 1.0);
    private_nh_.param("ndt_step_size", ndt_step_size_, 0.1);
    private_nh_.param("ndt_epsilon", ndt_epsilon_, 0.01);
    private_nh_.param("max_iterations", max_iterations_, 30);
    
    // 坐标系参数
    private_nh_.param<std::string>("map_frame_id", map_frame_id_, "map");
    private_nh_.param<std::string>("base_frame_id", base_frame_id_, "base_link");
    private_nh_.param<std::string>("lidar_frame_id", lidar_frame_id_, "velodyne");
}

void NDTMapper::pointCloudCallback(const sensor_msgs::PointCloud2::ConstPtr& input_cloud) {
    pcl::PointCloud<pcl::PointXYZ>::Ptr raw_cloud(new pcl::PointCloud<pcl::PointXYZ>);
    pcl::fromROSMsg(*input_cloud, *raw_cloud);
    
    // 对输入点云进行降采样
    downsamplePointCloud(raw_cloud, current_cloud_, voxel_leaf_size_);
    
    if (is_first_frame_) {
        // 如果是第一帧，直接作为地图
        *map_cloud_ = *current_cloud_;
        current_pose_ = Eigen::Matrix4f::Identity();
        previous_pose_ = current_pose_;
        is_first_frame_ = false;
        
        // 发布地图和位姿
        publishMap();
        publishPose();
        publishCurrentScan();
        return;
    }
    
    // 获取初始位姿估计
    Eigen::Matrix4f initial_guess;
    if (!getInitialGuess(initial_guess)) {
        ROS_WARN("Failed to get initial guess, using previous pose");
        initial_guess = previous_pose_;
    }
    
    // 执行点云配准
    ndt_.setInputTarget(map_cloud_);
    ndt_.setInputSource(current_cloud_);
    
    pcl::PointCloud<pcl::PointXYZ>::Ptr output_cloud(new pcl::PointCloud<pcl::PointXYZ>);
    ndt_.align(*output_cloud, initial_guess);
    
    if (ndt_.hasConverged()) {
        // 获取变换矩阵
        current_pose_ = ndt_.getFinalTransformation();
        
        // 更新地图
        updateMap();
        
        // 发布结果
        publishMap();
        publishPose();
        publishCurrentScan();
        
        // 更新previous_pose_
        previous_pose_ = current_pose_;
        
        ROS_INFO_THROTTLE(1.0, "NDT score: %f", ndt_.getFitnessScore());
    } else {
        ROS_WARN("NDT did not converge!");
    }
}

bool NDTMapper::getInitialGuess(Eigen::Matrix4f& initial_guess) {
    try {
        // 获取lidar到base_link的转换
        geometry_msgs::TransformStamped lidar_to_base = 
            tf_buffer_.lookupTransform(base_frame_id_, lidar_frame_id_, ros::Time(0));
            
        // 将TransformStamped转换为Eigen::Matrix4f
        Eigen::Affine3d affine;
        tf2::doTransform(Eigen::Affine3d::Identity(), affine, lidar_to_base);
        Eigen::Matrix4f lidar_to_base_mat = affine.matrix().cast<float>();
        
        // 使用上一帧的位姿作为初始估计
        initial_guess = previous_pose_ * lidar_to_base_mat;
        return true;
    } catch (tf2::TransformException& ex) {
        ROS_WARN("%s", ex.what());
        return false;
    }
}

void NDTMapper::updateMap() {
    // 将当前点云转换到地图坐标系
    pcl::PointCloud<pcl::PointXYZ>::Ptr transformed_cloud(new pcl::PointCloud<pcl::PointXYZ>);
    pcl::transformPointCloud(*current_cloud_, *transformed_cloud, current_pose_);
    
    // 将转换后的点云添加到地图中
    *map_cloud_ += *transformed_cloud;
    
    // 对地图进行降采样以控制点云大小
    downsamplePointCloud(map_cloud_, map_cloud_, voxel_leaf_size_);
}

void NDTMapper::publishMap() {
    sensor_msgs::PointCloud2 map_msg;
    pcl::toROSMsg(*map_cloud_, map_msg);
    map_msg.header.frame_id = map_frame_id_;
    map_msg.header.stamp = ros::Time::now();
    map_pub_.publish(map_msg);
}

void NDTMapper::publishCurrentScan() {
    sensor_msgs::PointCloud2 scan_msg;
    pcl::toROSMsg(*current_cloud_, scan_msg);
    scan_msg.header.frame_id = lidar_frame_id_;
    scan_msg.header.stamp = ros::Time::now();
    current_scan_pub_.publish(scan_msg);
}

void NDTMapper::publishPose() {
    geometry_msgs::PoseStamped pose_msg;
    pose_msg.header.frame_id = map_frame_id_;
    pose_msg.header.stamp = ros::Time::now();
    
    // 从变换矩阵中提取位置和姿态
    Eigen::Vector3f translation = current_pose_.block<3,1>(0,3);
    Eigen::Matrix3f rotation = current_pose_.block<3,3>(0,0);
    Eigen::Quaternionf q(rotation);
    
    pose_msg.pose.position.x = translation.x();
    pose_msg.pose.position.y = translation.y();
    pose_msg.pose.position.z = translation.z();
    pose_msg.pose.orientation.x = q.x();
    pose_msg.pose.orientation.y = q.y();
    pose_msg.pose.orientation.z = q.z();
    pose_msg.pose.orientation.w = q.w();
    
    current_pose_pub_.publish(pose_msg);
    
    // 发布TF变换
    geometry_msgs::TransformStamped transform_stamped;
    transform_stamped.header = pose_msg.header;
    transform_stamped.child_frame_id = base_frame_id_;
    transform_stamped.transform.translation.x = translation.x();
    transform_stamped.transform.translation.y = translation.y();
    transform_stamped.transform.translation.z = translation.z();
    transform_stamped.transform.rotation = pose_msg.pose.orientation;
    
    tf_broadcaster_.sendTransform(transform_stamped);
}

void NDTMapper::downsamplePointCloud(const pcl::PointCloud<pcl::PointXYZ>::Ptr& input_cloud,
                                   pcl::PointCloud<pcl::PointXYZ>::Ptr& output_cloud,
                                   double leaf_size) {
    pcl::VoxelGrid<pcl::PointXYZ> voxel_grid;
    voxel_grid.setLeafSize(leaf_size, leaf_size, leaf_size);
    voxel_grid.setInputCloud(input_cloud);
    voxel_grid.filter(*output_cloud);
}

bool NDTMapper::saveMap(const std::string& filename) {
    if (!map_cloud_ || map_cloud_->empty()) {
        ROS_ERROR("No map data to save!");
        return false;
    }

    // 保存地图到PCD文件
    if (pcl::io::savePCDFileBinary(filename, *map_cloud_) == -1) {
        ROS_ERROR("Failed to save map to %s", filename.c_str());
        return false;
    }
    
    ROS_INFO("Successfully saved map to %s", filename.c_str());
    return true;
}

bool NDTMapper::saveMapCallback(std_srvs::Trigger::Request& req, std_srvs::Trigger::Response& res) {
    std::string filename = ros::package::getPath("auto_navigation") + "/maps/map_" + 
                          std::to_string(ros::Time::now().toSec()) + ".pcd";
    if (saveMap(filename)) {
        res.success = true;
        res.message = "Map saved to " + filename;
        ROS_INFO_STREAM(res.message);
    } else {
        res.success = false;
        res.message = "Failed to save map";
        ROS_ERROR_STREAM(res.message);
    }
    return true;
}

} // namespace auto_navigation 